import pefile
import os
import sys
import re
import argparse
from collections import defaultdict

def demangle_cpp_name(mangled_name):
    """
    简单解修饰C++名称（针对MSVC风格的修饰）
    更复杂的解修饰可能需要调用外部工具如undname.exe
    """
    # MSVC风格的C++修饰名称通常以'?'开头
    if not mangled_name.startswith('?'):
        return None

    # 简单处理常见模式（实际情况更复杂）
    patterns = [
        # 匹配类似 "?functionName@namespace@@..." 的模式
        (r'\?(\w+)@(\w+)@@.*', r'\2::\1'),
        # 匹配简单函数名
        (r'\?(\w+)@@.*', r'\1'),
        # 处理带有参数的情况
        (r'\?(\w+)@.*?@Z', r'\1'),
    ]

    for pattern, replacement in patterns:
        demangled = re.sub(pattern, replacement, mangled_name)
        if demangled != mangled_name:
            return demangled

    return mangled_name

def analyze_dll(dll_path):
    """分析DLL文件的依赖和导出函数"""
    if not os.path.exists(dll_path):
        print(f"错误: 文件 '{dll_path}' 不存在")
        return None

    try:
        pe = pefile.PE(dll_path)
    except Exception as e:
        print(f"解析DLL失败: {str(e)}")
        return None

    result = {
        'path': dll_path,
        'filename': os.path.basename(dll_path),
        'dependencies': {
            'required': {},
            'missing': []
        },
        'exports': {
            'raw_functions':[],
            'functions': [],
            'variables': defaultdict(list)
        }
    }

    # 分析依赖项
    if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
        for entry in pe.DIRECTORY_ENTRY_IMPORT:
            dll_name = entry.dll.decode('utf-8')
            result['dependencies']['required'][dll_name]={'raw_functions':[],'functions': []}
            
            for exp in entry.imports:
                if exp.name:
                    raw_func_name = exp.name.decode('utf-8')
                    result['dependencies']['required'][dll_name]['raw_functions'].append(raw_func_name)
                    if re.findall(r'[^a-zA-Z0-9_]',raw_func_name):
                        func_name = demangle_cpp_name(raw_func_name)
                        result['dependencies']['required'][dll_name]['functions'].append(func_name)
            
            # 简单检查系统目录中是否存在该DLL
            system_dirs = os.environ.get('Path', '').split(';')
            system_dirs.append(os.path.dirname(dll_path))
            found = False
            for sys_dir in system_dirs:
                if os.path.exists(os.path.join(sys_dir, dll_name)):
                    found = True
                    break
            
            if not found:
                del result['dependencies']['required'][dll_name]
                result['dependencies']['missing'].append(dll_name)

    # 分析导出函数
    if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT'):
        exports = pe.DIRECTORY_ENTRY_EXPORT.symbols
        for exp in exports:
            if exp.name:
                raw_func_name = exp.name.decode('utf-8')
                result['exports']['raw_functions'].append(raw_func_name)
                if re.findall(r'[^a-zA-Z0-9_]',raw_func_name):
                    func_name = demangle_cpp_name(raw_func_name)
                    result['exports']['functions'].append(func_name)
                result['exports']['variables'][exp.ordinal].append(func_name)

    pe.close()
    return result

def print_analysis(result,args):
    """打印分析结果"""
    if not result:
        return

    print(f"\nDLL分析结果: {result['filename']}")
    print("=" * 50)

    # 打印依赖信息
    if args.depends:
        print("\n依赖项信息:")
        
        if result['dependencies']['required']:
            print(f"  找到依赖数: {len(result['dependencies']['required'])}")
            for founding in result['dependencies']['required']:
                print(f"    - {founding}")
                if args.imports:
                    for symbols in result['dependencies']['required'][founding]['functions']:
                        print(f'            +{symbols}')

        if result['dependencies']['missing']:
            print(f"  缺失依赖: {len(result['dependencies']['missing'])}")
            for missing in result['dependencies']['missing']:
                print(f"    - {missing}")
    
    # 打印导出函数信息
    if args.symbols:
        print("\n导出函数信息:")
        print(f"  导出函数总数: {len(result['exports']['raw_functions'])}")
        for func in result['exports']['raw_functions']:
            print(f"    - {func}")

    if args.code_symbols:
        print("\n导出c++函数信息:")
        print(f"  导出c++函数总数: {len(result['exports']['functions'])}")
        for func in result['exports']['functions']:
            print(f"    - {func}")

    if args.variables:
        print("\n导出变量:")
        print(f"  导出变量总数: {len(result['exports']['variables'])}")
        for o in result['exports']['variables']:
            print(f"    - {o}")

    print("\n" + "=" * 50)

if __name__ == "__main__":
    parse = argparse.ArgumentParser('command help')
    parse.add_argument("dll",help="input a dll file")
    parse.add_argument('-i','--imports',action='store_true',default=None)
    parse.add_argument('-v','--variables',action='store_true',default=None)
    parse.add_argument('-s','--symbols',action='store_true',default=None)
    parse.add_argument('-S','--code_symbols',action='store_true',default=None)
    parse.add_argument('-d','--depends',action='store_true',default=None)

    args = parse.parse_args()
    analysis_result = analyze_dll(args.dll)
    print_analysis(analysis_result,args)
    